from typing import List

import numpy as np

import torch
import torch.nn as nn
import librosa
from torchvision import transforms as transforms

try:
    from .utils import compute_dB
except:
    from utils import compute_dB


class Mel_Spectrogram(object):
    def __init__(self, sample_rate=16000, n_fft=512, win_length=400, hop=160, n_mels=80):
        super(Mel_Spectrogram, self).__init__()
        self.n_fft = n_fft
        self.n_mels = n_mels
        self.win_length = win_length
        self.hop = hop

        mel_basis = librosa.filters.mel(
            sr=sample_rate, n_fft=n_fft, n_mels=n_mels)
        self.mel_basis = torch.FloatTensor(mel_basis).cuda()
        window = torch.hamming_window(self.win_length)
        self.window = torch.FloatTensor(window).cuda()
        self.instance_norm = nn.InstanceNorm1d(n_mels)

    def __call__(self, x):
        x = torch.stft(x, n_fft=self.n_fft, hop_length=self.hop,
                       window=self.window, win_length=self.win_length, return_complex=True)
        x = torch.abs(x)
        x += 1e-4
        x = torch.log(x)
        x = torch.matmul(self.mel_basis, x)
        x = self.instance_norm(x)
        return x.unsqueeze(1)


class Add_Gaussian_Noise(object):
    def __init__(self, snr=[20, 30]):
        self.snr = snr

    def __call__(self, waveform):
        """
        Args:
            x (torch.tensor): Input waveform tensor (batch_size, waveform_length).
        Returns:
            torch.tensor: Output waveform tensor (batch_size, waveform_length).
        """
        snr = np.random.randint(self.snr[0], self.snr[1])
        clean_dB = compute_dB(waveform)
        noise = torch.randn(waveform.shape[-1]).cuda()
        noise_dB = compute_dB(noise)
        noise = torch.sqrt(10 ** ((clean_dB - noise_dB - snr) / 10)) * noise
        waveform = (waveform + noise)
        return waveform


class Spec_Crop(object):
    def __init__(self, frames=200):
        self.frames = frames 

    def __call__(self, spec):
        """
        Args:
            x (torch.tensor): Input Spec tensor (batch_size, n_mels, frame_length).
        Returns:
            torch.tensor: Output Spec tensor (batch_size, n_mels, frame_length).
        """
        return spec[:,:,:200]


class Spec_Swap(object):
    def __init__(self, frames=200):
        self.frames = frames 

    def __call__(self, spec):
        """
        Args:
            x (torch.tensor): Input spec Tensor (batch_size, n_mels, frame_length).
        Returns:
            torch.tensor: Output spec Tensor (batch_size, n_mels, frame_length).
        """
        return spec[:,:,:200]


class Train_Transform(object):
    def __init__(
        self,
        size_crops: List[int] = [96, 36],
        nmb_crops: List[int] = [2, 4],
    ):

        # 1. basic transform
        basic_trans = Mel_Spectrogram(n_mels=80)

        # 2. noise transform
        noise_trans = transforms.Compose([Add_Gaussian_Noise(), 
                basic_trans])

        # 3. crop transform
        crop_trans = transforms.Compose([
                basic_trans,
                Spec_Crop()
                ])

        # 4. noise crop transform
        noise_crop_trans = transforms.Compose([
                noise_trans,
                Spec_Crop()
                ])

        # 5. swap transform
        # swap_trans = transforms.Compose([
        #         basic_trans,
        #         Spec_Swap()
        #         ])

        # # 6. noise swap transform
        # noise_swap_trans = transforms.Compose([
        #         basic_trans,
        #         Spec_Swap()
        #         ])

        self.transform = []
        self.transform.append(basic_trans)
        self.transform.append(noise_trans)
        self.transform.append(crop_trans)
        self.transform.append(noise_crop_trans)
        #self.transform.append(swap_trans)
        #self.transform.append(noise_swap_trans)

    def __call__(self, waveform):
        multi_crops = list(map(lambda transform: transform(waveform), self.transform))
        return multi_crops


if __name__ == "__main__":
    from scipy.io import wavfile

    sample_rate, sig = wavfile.read("test.wav")
    sig = sig[:16000*3]
    sig = torch.FloatTensor(sig.copy())
    waveform = sig.repeat(128, 1).cuda()

    trans = Train_Transform()
    feature = trans(waveform)

    print(feature[0].shape)
    print(feature[1].shape)
    print(feature[2].shape)
    print(feature[3].shape)
    #print(feature[4].shape)
    #print(feature[5].shape)

    import matplotlib.pyplot as plt

    plt.subplot(411)
    plt.imshow(feature[0][0][0].cpu())

    plt.subplot(412)
    plt.imshow(feature[1][0][0].cpu())

    plt.subplot(413)
    plt.imshow(feature[2][0][0].cpu())

    plt.subplot(414)
    plt.imshow(feature[3][0][0].cpu())

    plt.tight_layout()

    plt.show()
    plt.savefig("test.png")

